#include "yolov8_detect.h"

#include <iostream> // 添加调试输出

Yolov8Detect::Yolov8Detect()
{
}

Yolov8Detect::~Yolov8Detect()
{
}

std::vector<DetectBox> Yolov8Detect::Inference(const cv::Mat input_mat, float scale,
                                               std::vector<int> &padding,
                                               float score_threshold,
                                               float iou_threshold,
                                               float nms_threshold)
{
    std::vector<DetectBox> detect_boxes;
    int channels = input_mat.channels();
    int height = input_mat.rows;
    int width = input_mat.cols;

    std::vector<float> input_data;
    input_data.resize(channels * height * width);
    // 归一化
    for (int c = 0; c < channels; c++)
    {
        for (int h = 0; h < height; h++)
        {
            for (int w = 0; w < width; w++)
            {
                float pix = input_mat.ptr<uchar>(h)[w * 3 + c];
                int chw_index = c * height * width + h * width + w;
                float tmp = pix / 255.0f;
                // input_data[c * height * width + h * width + w] = input_mat.at<cv::Vec3b>(h, w)[c];
                input_data[chw_index] = tmp;
            }
        }
    }

    // 创建输入tensor
    Ort::MemoryInfo memory_info = Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault);
    std::vector<int64_t> input_node_dims = {1, channels, height, width};

    Ort::Value input_tensor = Ort::Value::CreateTensor<float>(
        memory_info,
        input_data.data(),
        input_data.size(),
        input_node_dims.data(),
        input_node_dims.size());

    assert(input_tensor.IsTensor());

    // 获取输入节点名称
    std::vector<const char *> inputNodeNames = {"images"};
    std::vector<const char *> outputNodeNames = {"output0"};

    // onnx推理
    auto output_tensors = m_session->Run(
        Ort::RunOptions{nullptr},
        inputNodeNames.data(),
        &input_tensor,
        1,
        outputNodeNames.data(),
        outputNodeNames.size());

    // 获取输出维度
    // Ort::Value &predictions = output_tensors.at(0);
    // auto pred_dims = predictions.GetTensorTypeAndShapeInfo().GetShape();
    // std::cout << "Predictions dims: ";
    // for (auto i : pred_dims)
    // {
    //     std::cout << i << " ";
    // }
    // std::cout << "\n";

    // 后处理

    std::vector<int> class_ids;
    std::vector<float> confidences;
    std::vector<cv::Rect> boxes;

    auto output_data = output_tensors[0].GetTensorMutableData<float>(); // 获取输出数据

    // 输出维度
    auto output_dims = output_tensors[0].GetTensorTypeAndShapeInfo().GetShape();

    int64_t rows = output_dims[1]; // 10
    int64_t clos = output_dims[2]; // 8400

    // [10, 8400] -> [8400, 10]
    cv::Mat result = cv::Mat(rows, clos, CV_32F, output_data).t();
    // 数据宽度
    int new_width = class_names.size() + 4;

    float *pdata = (float *)result.data;

    for (int i = 0; i < clos; i++)
    {
        // 获取类别得分
        cv::Mat scores(1, class_names.size(), CV_32F, pdata + 4);
        cv::Point classIdPoint;
        double score;
        // 每一行得分的最大值
        cv::minMaxLoc(scores, 0, &score, 0, &classIdPoint);
        score = (float)score;
        if (score > score_threshold)
        {
            // 保存得分
            confidences.push_back(score);
            // 保存类别
            class_ids.push_back(classIdPoint.x);

            // 提取矩形框信息
            float x = pdata[0];
            float y = pdata[1];
            float w = pdata[2];
            float h = pdata[3];
            int left = int((x - padding[0] - 0.5 * w) * scale);
            int top = int((y - padding[1] - 0.5 * h) * scale);
            int width = int(w * scale);
            int height = int(h * scale);

            boxes.push_back(cv::Rect(left, top, width, height));
            // std::cout << "left: " << left << std::endl;
            // std::cout << "confidence: " << confidence << std::endl;
        }
        // 指向下一个数据
        pdata += new_width;
    }
    // nms
    std::vector<int> nmsResult;
    cv::dnn::NMSBoxes(boxes, confidences, score_threshold, iou_threshold, nmsResult);

    for (int i = 0; i < nmsResult.size(); i++)
    {
        DetectBox box;
        // 获取索引
        int idx = nmsResult[i];
        // 获取框
        box.left = boxes[idx].x;
        box.top = boxes[idx].y;
        box.right = boxes[idx].x + boxes[idx].width;
        box.bottom = boxes[idx].y + boxes[idx].height;
        // 获取置信度
        box.score = confidences[idx];
        // 标签id
        box.label = class_ids[idx];
        // 标签名
        box.label_name = class_names[class_ids[idx]];
        detect_boxes.push_back(box);
    }

    return detect_boxes;
}
